import base64
import requests
from PIL import Image
from io import BytesIO
from urllib.parse import urlparse, unquote
from mq.client_zmq import ZeroMQClient
from mq.client_redis import RedisClient
from abc import ABC, abstractmethod

class BaseWorker(ABC):
    def __init__(self, algo_type, mq_type, address):
        if mq_type == 'zmq':
            self.client = ZeroMQClient(address, algo_type, self)
        else:
            self.client = RedisClient(address, algo_type, self)
        self.device = "cuda"

    def start(self):
        self.client.start()

    @abstractmethod
    def process(self, message):
        pass

    def get_image_url(self, image_path):
        """
        根据 image_path 的类型(file://、URL 或 Base64)构造请求 URL
        """
        if image_path.startswith("file://"):
            # 如果是 file:// 开头的路径，解析出实际路径并转换为 Node.js 服务的 URL
            parsed_path = urlparse(image_path)
            actual_path = unquote(parsed_path.path)  # 解码路径
            return f"{self.node_server_address}/LoadImage?image_path={actual_path}"
        elif image_path.startswith("http://") or image_path.startswith("https://"):
            # 如果是 URL，直接返回
            return image_path
        elif image_path.startswith("data:image/"):
            # 如果是 Base64 编码的图片，直接返回
            return image_path
        else:
            raise ValueError("Unsupported image path format")

    def fetch_image(self, image_path):
        """
        根据 image_path 获取图片并转换为 NumPy 数组
        """
        # 获取图片的 URL 或 Base64 数据
        image_url_or_data = self.get_image_url(image_path)
        if image_url_or_data.startswith("data:image/"):
            # 如果是 Base64 编码的图片
            try:
                # 提取 Base64 数据部分（去掉 data:image/...;base64, 前缀）
                base64_index = image_url_or_data.find(";base64,")
                if base64_index == -1:
                    raise ValueError("Invalid Base64 image format")

                base64_content = image_url_or_data[base64_index + len(";base64,"):]  # 获取 Base64 数据
                # 解码 Base64 数据为二进制
                image_data = base64.b64decode(base64_content)
                # 将二进制数据转换为 PIL Image 对象
                image = Image.open(BytesIO(image_data))
                # 将图片转换为 RGB 格式并转换为 NumPy 数组
                return image.convert("RGB")
            except Exception as e:
                raise Exception(f"Failed to decode Base64 image: {str(e)}")
        else:
            # 如果是 URL 或本地文件路径
            try:
                # 发送 GET 请求获取图片
                response = requests.get(image_url_or_data)
                # 检查请求是否成功
                if response.status_code == 200:
                    # 将响应内容转换为 PIL Image 对象
                    image = Image.open(BytesIO(response.content))
                    # 将图片转换为 RGB 格式并转换为 NumPy 数组
                    return image.convert("RGB")
                else:
                    raise Exception(f"Failed to fetch image: {response.status_code}")
            except Exception as e:
                raise Exception(f"Failed to fetch image: {str(e)}")

    def truncate_message(self, message, max_length=50):
        """
        截断 message 中 inputs.image_path 字段的 Base64 编码字符串。

        :param message: 包含 inputs 字段的字典
        :param max_length: 截断后的最大长度，默认为 50
        :return: 处理后的 message
        """
        image_path = message['inputs']['image_path']
        if image_path and isinstance(image_path, str) and image_path.startswith('data:image/'):
            # 浅拷贝 message，避免修改原始数据
            truncate_message = {**message}
            # 浅拷贝 inputs 字段
            truncate_input = {**message.get('inputs', {})}

            # 如果 image_path 存在且是 Base64 编码的字符串，则截断
            truncate_input['image_path'] = image_path[:max_length] + '...'

            # 更新 inputs 字段
            truncate_message['inputs'] = truncate_input
            return truncate_message
        else:
            return message
        